#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Python version: 3.6

import cv2
import torch
from torch import nn
from torch.utils.data import DataLoader, Dataset


def clean_dataset(dataset, idxs):
    dataset_clean = []
    clean_num = 0
    # 遍历原始数据集中的每个样本
    for i, data in enumerate(dataset):  # 获取图像和标签
        img = data[0]
        label = data[1]
        # 如果当前样本的索引在idxs中，则对图像进行预处理
        if i in idxs:
            # 对图像进行预处理
            img_np = img.cpu().numpy()  # 将tensor转换为numpy数组
            img_blur = cv2.GaussianBlur(img_np, (3, 3), 0)  # 对numpy数组进行高斯模糊
            img_blur_tensor = torch.from_numpy(img_blur)  # 将numpy数组转换为tensor
            # 将处理后的图像和标签加入新列表中
            dataset_clean.append((img_blur_tensor, label))
            clean_num = clean_num + 1
        else:
            dataset_clean.append((img, label))
        # 打印完成信息
    print("完成本地数据清洗,共清洗数据：", clean_num)
    # 返回处理后的数据集
    return dataset_clean


class DatasetSplit(Dataset):
    def __init__(self, dataset, idxs):
        self.dataset = dataset
        self.idxs = list(idxs)

    def __len__(self):
        return len(self.idxs)

    def __getitem__(self, item):
        image, label = self.dataset[self.idxs[item]]
        return image, label


class LocalUpdate(object):
    def __init__(self, args, dataset=None, idxs=None):
        self.args = args
        self.loss_func = nn.CrossEntropyLoss()
        self.selected_clients = []
        # 数据清洗
        dataset_clean = clean_dataset(dataset, idxs)
        self.ldr_train = DataLoader(DatasetSplit(dataset_clean, idxs), batch_size=self.args.local_bs, shuffle=True)

    def train(self, net):
        # 将神经网络设置为训练模式
        net.train()
        # train and update
        # 创建一个随机梯度下降优化器（SGD），用于更新神经网络中的参数。其中，lr 和 momentum 是优化器的超参数，需要在调用该函数时指定
        optimizer = torch.optim.SGD(net.parameters(), lr=self.args.lr, momentum=self.args.momentum)
        # 在每个本地训练轮次（args.local_ep）中，遍历训练数据集中的所有批次（ldr_train），并对每个批次进行训练
        epoch_loss = []
        for iter in range(self.args.local_ep):
            batch_loss = []
            for batch_idx, (images, labels) in enumerate(self.ldr_train):
                # 对于每个批次，首先将输入数据（images 和 labels）移动到指定的设备上（args.device）
                images, labels = images.to(self.args.device), labels.to(self.args.device)
                # 然后将神经网络的梯度清零（net.zero_grad()）
                net.zero_grad()
                # 使用神经网络对输入数据进行前向传播得到预测值（log_probs）
                log_probs = net(images)
                # # 再计算预测值和真实标签之间的损失值（loss）
                loss = self.loss_func(log_probs, labels)
                # 通过反向传播计算梯度并更新神经网络的参数（loss.backward() 和 optimizer.step()）
                loss.backward()
                optimizer.step()
                # 如果 args.verbose 参数为 True，则在每个批次结束时打印出当前轮次、批次、总数据量、完成百分比和损失值等信息
                if self.args.verbose and batch_idx % 10 == 0:
                    print('Update Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                        iter, batch_idx * len(images), len(self.ldr_train.dataset),
                              100. * batch_idx / len(self.ldr_train), loss.item()))
                # 将当前轮次的每个批次的损失值记录下来，并求出平均损失作为该轮次的总体损失。将每个本地训练轮次的平均损失保存在 epoch_loss 列表中
                batch_loss.append(loss.item())
            epoch_loss.append(sum(batch_loss) / len(batch_loss))
        # 释放之前加载到内存中的数据集
        del self.ldr_train
        # 返回训练后的神经网络参数（net.state_dict()）和整个训练过程中的平均损失（sum(epoch_loss) / len(epoch_loss)）
        return net.state_dict(), sum(epoch_loss) / len(epoch_loss)
